# Copyright (C) 2022-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).

# --------------------------------------------------------
# Main encoder/decoder blocks
# --------------------------------------------------------
# References:
# timm
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/helpers.py
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/drop.py
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/mlp.py
# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/layers/patch_embed.py

import torch
import torch.nn as nn

from itertools import repeat
import collections.abc

# 返回重复n个x的元组
def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): # x是非str的可迭代对象
            return x
        return tuple(repeat(x, n)) # 返回有n个x的元组
    return parse
to_2tuple = _ntuple(2) # 返回有2个x的元组

def drop_path(x, drop_prob: float = 0., training: bool = False, scale_by_keep: bool = True):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor

# DropPath正则化，随机删除多分支结构
class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
    """
    def __init__(self, drop_prob: float = 0., scale_by_keep: bool = True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

    def extra_repr(self):
        return f'drop_prob={round(self.drop_prob,3):0.3f}'

# MLP全连接层
class Mlp(nn.Module):
    """ MLP as used in Vision Transformer, MLP-Mixer and related networks"""
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, bias=True, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)
        drop_probs = to_2tuple(drop)

        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0]) # 线性层
        self.act = act_layer() # 默认GELU
        self.drop1 = nn.Dropout(drop_probs[0]) # Dropout正则化
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
        self.drop2 = nn.Dropout(drop_probs[1])

    def forward(self, x):
        x = self.fc1(x) # 线性层
        x = self.act(x) # GELU
        x = self.drop1(x) # Dropout正则化
        x = self.fc2(x) # 线性层
        x = self.drop2(x) # Dropout正则化
        return x

# 自注意力
class Attention(nn.Module):
    def __init__(self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads # 多头数量
        head_dim = dim // num_heads # 每个头的维度(均分维度的策略)
        self.scale = head_dim ** -0.5 # 注意力的尺度缩放(sqrt(dk))

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) # 线性层，输出维度*3
        self.attn_drop = nn.Dropout(attn_drop) # Dropout正则化
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.rope = rope # 旋转位置编码

    def forward(self, x, xpos):
        B, N, C = x.shape

        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).transpose(1,3) # 使用线性层维度*3获得qkv，改变形状(B,N,头数量,头维度)，交换维度(B,头数量,N,头维度)
        q, k, v = [qkv[:,:,i] for i in range(3)] # 沿第2维度从qkv张量中分出q(query)、k(key)和v(value)，(B,头数量,N,头维度)
        # q,k,v = qkv.unbind(2)  # make torchscript happy (cannot use tensor as tuple)

        if self.rope is not None: # q和k加入旋转位置编码
            q = self.rope(q, xpos) # 输入q，补丁位置编码，输出加入旋转位置编码后的q
            k = self.rope(k, xpos)

        attn = (q @ k.transpose(-2, -1)) * self.scale # 注意力张量，(Q@K^T)/sqrt(dk)，@是矩阵乘法
        attn = attn.softmax(dim=-1) # 注意力权重，在嵌入(头)维度做softmax
        attn = self.attn_drop(attn) # Dropout正则化

        x = (attn @ v).transpose(1, 2).reshape(B, N, C) # 缩放点积注意力，注意力权重@v，合并所有头
        x = self.proj(x) # 全连接
        x = self.proj_drop(x) # Dropout正则化
        return x

# 编码器块
class Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, rope=None):
        super().__init__()
        self.norm1 = norm_layer(dim) # 层归一化
        self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) # 自注意力
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() # DropPath正则化，随机删除多分支结构
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio) # 扩展MLP全连接层的隐藏层维度
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) # MLP全连接层

    def forward(self, x, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos)) # 自注意力(q和k加入旋转位置编码)，DropPath正则化，残差连接
        x = x + self.drop_path(self.mlp(self.norm2(x))) # 层归一化，MLP全连接层，DropPath正则化，残差连接
        return x

# 交叉注意力
class CrossAttention(nn.Module):
    def __init__(self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads # 多头数量
        head_dim = dim // num_heads # 每个头的维度(均分维度的策略)
        self.scale = head_dim ** -0.5 # 注意力的尺度缩放(sqrt(dk))

        self.projq = nn.Linear(dim, dim, bias=qkv_bias) # 线性层，维度不变
        self.projk = nn.Linear(dim, dim, bias=qkv_bias)
        self.projv = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop) # Dropout正则化
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.rope = rope # 旋转位置编码

    def forward(self, query, key, value, qpos, kpos): # query=token特征1，key=value=token特征2
        B, Nq, C = query.shape
        Nk = key.shape[1] # key的嵌入长度
        Nv = value.shape[1] # value的嵌入长度

        q = self.projq(query).reshape(B,Nq,self.num_heads,C//self.num_heads).permute(0, 2, 1, 3) # 使用线性层获得q，改变形状(B,N,头数量,头维度)，交换维度(B,头数量,N,头维度)
        k = self.projk(key).reshape(B,Nk,self.num_heads,C//self.num_heads).permute(0, 2, 1, 3)
        v = self.projv(value).reshape(B,Nv,self.num_heads,C//self.num_heads).permute(0, 2, 1, 3)

        if self.rope is not None: # q和k加入旋转位置编码
            q = self.rope(q, qpos) # 输入q，补丁位置编码，输出加入旋转位置编码后的q
            k = self.rope(k, kpos)

        attn = (q @ k.transpose(-2, -1)) * self.scale # 注意力，(Q*K^T)/sqrt(dk)
        attn = attn.softmax(dim=-1) # 注意力权重，在嵌入(头)维度做softmax
        attn = self.attn_drop(attn) # Dropout正则化

        x = (attn @ v).transpose(1, 2).reshape(B, Nq, C) # 注意力权重叉乘(@)v，合并所有头
        x = self.proj(x) # 全连接
        x = self.proj_drop(x) # Dropout正则化
        return x

# 解码器块
class DecoderBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, drop=0., attn_drop=0.,
                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, norm_mem=True, rope=None):
        super().__init__()
        self.norm1 = norm_layer(dim) # 层归一化
        self.attn = Attention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) # 自注意力
        self.cross_attn = CrossAttention(dim, rope=rope, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) # 交叉注意力(k和v来自第2视图)
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() # DropPath正则化，随机删除多分支结构
        self.norm2 = norm_layer(dim)
        self.norm3 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio) # 扩展MLP全连接层的隐藏层维度
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) # MLP全连接层
        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity() # 依情况对token特征2做层归一化

    def forward(self, x, y, xpos, ypos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos)) # 自注意力(q和k加入旋转位置编码)，DropPath正则化，残差连接
        y_ = self.norm_y(y) # 依情况对token特征2做层归一化
        x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, y_, xpos, ypos))# 交叉注意力(q=层归一化token特征1，k=v=token特征2)，DropPath正则化，残差连接
        x = x + self.drop_path(self.mlp(self.norm3(x))) # 层归一化，MLP全连接层，DropPath正则化，残差连接
        return x, y # 其中，y(token特征2)没有改变，也没有再使用

# patch embedding，获取补丁的位置编码
class PositionGetter(object):
    """ return positions of patches """

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h,w) in self.cache_positions: # cache_positions字典中保存(h,w)尺寸的位置编码
            x = torch.arange(w, device=device) # 创建张量，[0,w-1]
            y = torch.arange(h, device=device)
            self.cache_positions[h,w] = torch.cartesian_prod(y, x) # (h, w, 2)，x和y的笛卡尔积
        pos = self.cache_positions[h,w].view(1, h*w, 2).expand(b, -1, 2).clone()
        return pos

# 补丁嵌入
class PatchEmbed(nn.Module):
    """ just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
        super().__init__()
        img_size = to_2tuple(img_size) # 图像尺寸，默认(224,224)，实际是元组不进行to_2tuple
        patch_size = to_2tuple(patch_size) # 补丁尺寸，(16,16)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1]) # 网格尺寸
        self.num_patches = self.grid_size[0] * self.grid_size[1] # 补丁数量
        self.flatten = flatten # 扁平化

        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) # 2d卷积，默认输入3通道，输出768通道，卷积核=补丁尺寸(16*16)，步长=补丁尺寸(16)
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() # 层归一化

        self.position_getter = PositionGetter() # 获取补丁的位置编码

    def forward(self, x):
        B, C, H, W = x.shape
        torch._assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).") # 输入图像尺寸需要与模型一致
        torch._assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
        x = self.proj(x) # 2d卷积，默认输入3通道，输出768通道，卷积核=补丁尺寸(16*16)，步长=补丁尺寸(16)
        pos = self.position_getter(B, x.size(2), x.size(3), x.device) # 补丁的位置编码，输入卷积后的B，H，W，device
        if self.flatten: # 扁平化
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BCN -> BNC
        x = self.norm(x) # 层归一化
        return x, pos # 返回token和位置编码

    # 权重初始化
    def _init_weights(self):
        w = self.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
